#include <iostream>
#include <PiccoloGeometryUtil.h>
#include <Graphics/PiccoloVKBuffer.h>
#include <Graphics/PiccoloVKGraphicsContext.h>
#include <Graphics/PiccoloVKImage.h>
#include <Graphics/PiccoloVKQueue.h>

#include "PiccoloLog.h"
#include "PiccoloWindow.h"
#include "Graphics/PiccoloGraphicsContext.h"
#include "Graphics/PiccoloVKDevice.h"
#include "Graphics/PiccoloVKSwapchain.h"
#include "Graphics/PiccoloVKRenderPass.h"
#include "Graphics/PiccoloVKFrameBuffer.h"
#include "Graphics/PiccoloVKPipeline.h"
#include "Graphics/PiccoloVKCommandBuffer.h"

#include "PiccoloFileUtil.h"

using namespace Piccolo;

struct PushConstants
{
    glm::mat4 matrix { 1.f };
};

int main()
{
    PiccoloLog::Init();

    std::unique_ptr<PiccoloWindow> window = PiccoloWindow::Create(800, 600, "FirstCube");
    std::unique_ptr<PiccoloGraphicsContext> context = PiccoloGraphicsContext::Create(window.get());
    auto* vkContext = dynamic_cast<PiccoloVKGraphicsContext*>(context.get());
    std::shared_ptr<PiccoloVKDevice> device = std::make_shared<PiccoloVKDevice>(vkContext, 1, 1);
    std::shared_ptr<PiccoloVKSwapchain> swapchain = std::make_shared<PiccoloVKSwapchain>(vkContext, device.get());
    swapchain->RecreateSwapchain();

    //TODO query depth format
    VkFormat depthFormat = VK_FORMAT_D32_SFLOAT;

    // 创建颜色附件和深度附件
    std::vector<VkAttachmentDescription> attachments =
    {
        {
            .flags = 0,
            .format = device->GetSettings().surfaceFormat,
            .samples = VK_SAMPLE_COUNT_1_BIT,
            .loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR,
            .storeOp = VK_ATTACHMENT_STORE_OP_STORE,
            .stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE,
            .stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE,
            .initialLayout = VK_IMAGE_LAYOUT_UNDEFINED,
            .finalLayout = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR
        },
        {
            .flags = 0,
            .format = depthFormat,
            .samples = VK_SAMPLE_COUNT_1_BIT,
            .loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR,
            .storeOp = VK_ATTACHMENT_STORE_OP_STORE,
            .stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE,
            .stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE,
            .initialLayout = VK_IMAGE_LAYOUT_UNDEFINED,
            .finalLayout = VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL
        }
    };
    std::vector<RenderSubPass> subPasses =
    {
        {
            .colorAttachment = { 0, VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL },
            .depthStencilAttachment = { 1, VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL }
        }
    };

    std::shared_ptr<PiccoloVKRenderPass> renderPass = std::make_shared<PiccoloVKRenderPass>(device.get(), attachments, subPasses);


    std::vector<VkImage> swapchainImages = swapchain->GetImages();
    uint32_t swapchainImageSize = swapchainImages.size();

    std::vector<std::shared_ptr<PiccoloVKFrameBuffer>> frameBuffers;

    VkExtent3D imageExtent = { swapchain->GetWidth(), swapchain->GetHeight(), 1 };

    for (int i = 0; i < swapchainImageSize; ++i)
    {
        std::vector<std::shared_ptr<PiccoloVKImage>> images =
        {
            std::make_shared<PiccoloVKImage>(device.get(), swapchainImages[i], imageExtent, device->GetSettings().surfaceFormat, VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT),
            std::make_shared<PiccoloVKImage>(device.get(), imageExtent, depthFormat, VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT)
        };
        frameBuffers.push_back(std::make_shared<PiccoloVKFrameBuffer>(device.get(), renderPass.get(), images, swapchain->GetWidth(), swapchain->GetHeight()));
    }

    std::vector<VkPushConstantRange> pushConstantRanges =
    {
        {
            .stageFlags = VK_SHADER_STAGE_VERTEX_BIT,
            .offset = 0,
            .size = sizeof(PushConstants)
        }
    };
    ShaderLayout shaderLayout =
    {
        .pushConstantRanges = pushConstantRanges
    };

    std::shared_ptr<PiccoloVKPipelineLayout> pipelineLayout = std::make_shared<PiccoloVKPipelineLayout>(device.get(),
        PICCOLO_RES_SHADER_DIR"cube.vert", PICCOLO_RES_SHADER_DIR"cube.frag", shaderLayout);
    std::shared_ptr<PiccoloVKPipeline> pipeline = std::make_shared<PiccoloVKPipeline>(device.get(), renderPass.get(), pipelineLayout.get());


    std::vector<VkVertexInputBindingDescription> vertexBindings =
    {
        {
            .binding = 0,
            .stride = sizeof(PiccoloVertex),
            .inputRate = VK_VERTEX_INPUT_RATE_VERTEX
        }
    };
    std::vector<VkVertexInputAttributeDescription> vertexAttributes =
    {
        {
            .location = 0,
            .binding = 0,
            .format = VK_FORMAT_R32G32B32_SFLOAT,
            .offset = offsetof(PiccoloVertex, position)
        },
        {
            .location = 1,
            .binding = 0,
            .format = VK_FORMAT_R32G32_SFLOAT,
            .offset = offsetof(PiccoloVertex, texcoord0)
        },
        {
            .location = 2,
            .binding = 0,
            .format = VK_FORMAT_R32G32B32_SFLOAT,
            .offset = offsetof(PiccoloVertex, normal)
        }
    };

    pipeline->SetVertexInputState(vertexBindings, vertexAttributes);
    pipeline->SetInputAssemblyState(VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST)->EnableDepthTest();
    pipeline->SetDynamicState({ VK_DYNAMIC_STATE_VIEWPORT, VK_DYNAMIC_STATE_SCISSOR });
    pipeline->Create();

    std::shared_ptr<PiccoloVKCommandPool> commandPool = std::make_shared<PiccoloVKCommandPool>(device.get(), vkContext->GetGraphicQueueFamilyInfo().queueFamilyIndex);
    std::vector<VkCommandBuffer> commandBuffers = commandPool->AllocateCommandBuffer(swapchainImages.size());

    // Geometry
    std::vector<PiccoloVertex> vertices;
    std::vector<uint32_t> indices;
    PiccoloGeometryUtil::CreateCube(-0.3f, 0.3f, -0.3f, 0.3f, -0.3f, 0.3f, vertices, indices);
    // 顶点缓冲区和索引缓冲区
    std::shared_ptr<PiccoloVKBuffer> vertexBuffer = std::make_shared<PiccoloVKBuffer>(device.get(), VK_BUFFER_USAGE_VERTEX_BUFFER_BIT, sizeof(vertices[0]) * vertices.size(), vertices.data());
    std::shared_ptr<PiccoloVKBuffer> indexBuffer = std::make_shared<PiccoloVKBuffer>(device.get(), VK_BUFFER_USAGE_INDEX_BUFFER_BIT, sizeof(indices[0]) * indices.size(), indices.data());

    PushConstants pushConstants {};

    PiccoloVKQueue* graphicQueue = device->GetFirstGraphicQueue();
    const std::vector<VkClearValue> clearValues =
    {
        {0.1f, 0.2f, 0.3f, 1.f},
        {1.f, 0.f}
    };


    // 获取交换链图像信号量
    // 提交信号量
    // frame fence

    const uint32_t numberOfBuffer = 2;
    std::vector<VkSemaphore> imageAvailableSemaphores(numberOfBuffer);
    std::vector<VkSemaphore> submittedSemaphores(numberOfBuffer);
    std::vector<VkFence> frameFences(numberOfBuffer);

    VkSemaphoreCreateInfo semaphoreInfo =
    {
        .sType = VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO,
        .pNext = nullptr,
        .flags = 0
    };
    VkFenceCreateInfo fenceInfo =
    {
        .sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO,
        .pNext = nullptr,
        .flags = VK_FENCE_CREATE_SIGNALED_BIT
    };
    for (int i = 0; i < numberOfBuffer; ++i)
    {
        CALL_VK(vkCreateSemaphore(device->GetLogicDevice(), &semaphoreInfo, nullptr, &imageAvailableSemaphores[i]));
        CALL_VK(vkCreateSemaphore(device->GetLogicDevice(), &semaphoreInfo, nullptr, &submittedSemaphores[i]));
        CALL_VK(vkCreateFence(device->GetLogicDevice(), &fenceInfo, nullptr, &frameFences[i]));
    }

    std::chrono::time_point lastTimePoint = std::chrono::steady_clock::now();

    uint32_t currentBuffer = 0;
    while (!window->ShouldClose())
    {
        window->PollEvents();

        CALL_VK(vkWaitForFences(device->GetLogicDevice(), 1, &frameFences[currentBuffer], VK_TRUE, UINT64_MAX));
        CALL_VK(vkResetFences(device->GetLogicDevice(), 1, &frameFences[currentBuffer]));

        // 从交换链获取image
        uint32_t imageIndex = swapchain->AcquireImage(imageAvailableSemaphores[currentBuffer]);

        float time = std::chrono::duration<float>(std::chrono::steady_clock::now() - lastTimePoint).count();
        pushConstants.matrix = glm::rotate(glm::mat4(1.f), glm::radians(-17.f), glm::vec3(1.f, 0.f, 0.f));
        pushConstants.matrix = glm::rotate(pushConstants.matrix, glm::radians(time * 100.f), glm::vec3(0.f, 1.f, 0.f));
        pushConstants.matrix = glm::ortho(-1.f, 1.f, -1.f, 1.f, -1.f, 1.f) * pushConstants.matrix;

        // 开启command buffer
        PiccoloVKCommandPool::BeginCommandBuffer(commandBuffers[imageIndex]);

        // 开启 render pass，绑定帧缓冲
        renderPass->BeginRenderPass(commandBuffers[imageIndex], frameBuffers[imageIndex].get(), clearValues);

        // 绑定资源（pipeline、geometry、descriptorSet）
        pipeline->Bind(commandBuffers[imageIndex]);

        VkViewport viewport =
        {
            .x = 0,
            .y = 0,
            .width = static_cast<float>(frameBuffers[imageIndex]->GetWidth()),
            .height = static_cast<float>(frameBuffers[imageIndex]->GetHeight()),
            .minDepth = 0,
            .maxDepth = 1
        };
        vkCmdSetViewport(commandBuffers[imageIndex], 0, 1, &viewport);
        VkRect2D scissor = { {0, 0}, {frameBuffers[imageIndex]->GetWidth(), frameBuffers[imageIndex]->GetHeight()} };
        vkCmdSetScissor(commandBuffers[imageIndex], 0, 1, &scissor);

        // 设置shader的常量缓冲区
        vkCmdPushConstants(commandBuffers[imageIndex], pipelineLayout->GetPipelineLayout(), VK_SHADER_STAGE_VERTEX_BIT, 0, sizeof(pushConstants), &pushConstants);

        // 绑定顶点缓冲区
        VkBuffer vertexBuffers[] = { vertexBuffer->GetBuffer() };
        VkDeviceSize offsets[] = { 0 };
        vkCmdBindVertexBuffers(commandBuffers[imageIndex], 0, 1, vertexBuffers, offsets);
        // 绑定索引缓冲区
        vkCmdBindIndexBuffer(commandBuffers[imageIndex], indexBuffer->GetBuffer(), 0, VK_INDEX_TYPE_UINT32);

        // 绘制
        vkCmdDrawIndexed(commandBuffers[imageIndex], indices.size(), 1, 0, 0, 0);
        // vkCmdDraw(commandBuffers[imageIndex], 3, 1, 0, 0);

        // 结束 render pass
        renderPass->EndRenderPass(commandBuffers[imageIndex]);

        // 结束 command buffer
        PiccoloVKCommandPool::EndCommandBuffer(commandBuffers[imageIndex]);

        // 提交 command buffer 到队列
        graphicQueue->Submit({ commandBuffers[imageIndex] }, { imageAvailableSemaphores[currentBuffer] }, { submittedSemaphores[currentBuffer] }, frameFences[currentBuffer]);
        // graphicQueue->WaitIdle();

        // 呈现 present
        swapchain->Present(imageIndex, { submittedSemaphores[currentBuffer] });

        window->SwapBuffer();

        currentBuffer = (currentBuffer + 1) % numberOfBuffer;
    }

    for (int i = 0; i < numberOfBuffer; ++i)
    {
        vkDeviceWaitIdle(device->GetLogicDevice());
        VK_D(Semaphore, device->GetLogicDevice(), imageAvailableSemaphores[i]);
        VK_D(Semaphore, device->GetLogicDevice(), submittedSemaphores[i]);
        VK_D(Fence, device->GetLogicDevice(), frameFences[i]);
    }

    return EXIT_SUCCESS;
}